import maya.cmds as mc


class LZ_ikspline():
    def __init__(self):
        pass
    def createIkSpline_baseOnCurve(self,curveInput = '',splineName = '', jntProportion = [1,1,1] , jntNum = 3 ,stretch = 1 ,keepvolume = 1 , stretchType = 'scale'):
        
        
        
        #-------------------------------------------------------------------------------    
        if len(jntProportion) == jntNum:
            
            allgrp = splineName + '_LZIkSpline' + '_grp'
            if mc.objExists(allgrp) == 0:
            
                mc.group(em=True,n= allgrp )
                mc.addAttr(allgrp,ln = 'stretch',k=True , dv = stretch  )
                mc.addAttr(allgrp,ln = 'globalScale',k=True , dv = stretch  )
                    
                proportionM = 0
                for num in range(jntNum):
                    proportionM = proportionM + jntProportion[num]
                    
                #-------------------------------------------------------------------------------
                curveShape = mc.listRelatives( curveInput , s=True )[0]
                
                
                #for jn in range(jntNum+1):
                #    
                #
                #    self.createSeketon(seketonInfo = [ ('' , (0,0,0) )] )

                #-------------------------------------------------------------------------------    
                distanNode_all = []
                position_all = []
                for i in range(jntNum):
                    curveInfo = splineName + '_LZIkSpline' + 'pointOnCurveInfo'+str(i)+'_poci'
                    mc.createNode('pointOnCurveInfo',n = curveInfo )
                    mc.setAttr(str( curveInfo )+'.turnOnPercentage' , 1)
                    mc.connectAttr(str(curveShape)+'.worldSpace' ,  str(curveInfo)+'.inputCurve')
                    
                    if i ==0:
                        currentParameter = 0
                    else:
                        currentParameter_value = 0
                        for  num_pro in range(i):
                            currentParameter_value = currentParameter_value + jntProportion[num]
                        
                        #print([currentParameter_value,proportionM])
                            
                        currentParameter = float(currentParameter_value)/proportionM
                        #print( [currentParameter_value,proportionM,currentParameter])
                    mc.setAttr( str(curveInfo)+'.parameter' , currentParameter )
                    position = mc.getAttr(str(curveInfo)+'.position')[0]
                    position_all.append(position)
                    
                    if i == jntNum -1:
                        curveInfoNext = splineName + '_LZIkSpline' + 'pointOnCurveInfo'+str(i+1)+'_poci'
                        mc.createNode('pointOnCurveInfo',n = curveInfoNext )
                        mc.connectAttr(str(curveShape)+'.worldSpace' ,  str(curveInfoNext)+'.inputCurve')
                        mc.setAttr(str( curveInfoNext )+'.turnOnPercentage' , 1)
                        mc.setAttr( str(curveInfoNext)+'.parameter' , 1 )
                        position = mc.getAttr(str(curveInfoNext)+'.position')[0]
                        position_all.append(position)
                    
                    
                    distanceNode = splineName + '_LZIkSpline' + 'distance' + str(i)+'_dis'      
                    mc.createNode('distanceBetween' , n = distanceNode )
                    
                    distanceNodeFront = splineName + '_LZIkSpline' + 'distance' + str(i-1)+'_dis'  
                    
                    mc.connectAttr( str(curveInfo) + '.position' ,str(distanceNode) + '.point1'  )
                    if i > 0:
                        mc.connectAttr( str(curveInfo) + '.position' ,str(distanceNodeFront) + '.point2'  )
                        
                    if i == jntNum-1:
                        mc.connectAttr( str(curveInfoNext) + '.position' ,str(distanceNode) + '.point2'  )
                    
                    
                    
                    
                    
                    distanNode_all.append(distanceNode)
                
                #-------------------------------------------------------------------------------
                curveShape = mc.listRelatives( curveInput , s=True )[0]
                #print(position_all)
                seketon_all = []

                for jn in range(jntNum+1):
                    
                    jnt_seketon = splineName + '_LZIkSpline' + '_seketon' + str(jn)+'_jnt'      
                    
                    if jn ==0:
                        secCancle = 1
                    else:
                        secCancle = 0
                    #print(jnt_seketon)
                    self.createSeketon(seketonInfo = [ (jnt_seketon , position_all[jn] )] , selectCancle = secCancle )  
                    
                    seketon_all.append( jnt_seketon ) 
                
                ikhandle_name =  splineName + '_LZIkSpline' + '_ikhandle'   
                mc.ikHandle( sj = seketon_all[0] , ee=seketon_all[-1] , sol = 'ikSplineSolver' , curve = curveInput, ccv = False ,n = ikhandle_name )    
                mc.parent( ikhandle_name , allgrp )
                mc.parent( seketon_all[0]  , allgrp )
                
                #-------------------------------------------------------------------------------    
                combine_md_all = []
                for num_pma in range(jntNum):
                    global_md = splineName + '_LZIkSpline' + '_multiply' + str(num_pma)+'_md'   
                    mc.createNode( 'multiplyDivide' , n = global_md )
                    orgDistance = mc.getAttr( str(distanNode_all[num_pma])+ '.distance'  )
                    mc.setAttr(str(global_md)+'.input1X' , orgDistance )
                    mc.connectAttr( str(allgrp)+'.globalScale' ,str(global_md) +'.input2X')
                
                    plusMinusAverage = splineName + '_LZIkSpline' + '_minus' + str(num_pma)+'_pma'   
                    mc.createNode(  'plusMinusAverage'  , n = plusMinusAverage )    
                    mc.setAttr(str(plusMinusAverage)+'.operation' , 2 )
                    mc.connectAttr( str(distanNode_all[num_pma])+ '.distance'  ,str(plusMinusAverage)+'.input1D[0]')
                    mc.connectAttr( str(global_md) +'.outputX'  ,str(plusMinusAverage)+'.input1D[1]')
                    
                    surplus_md = splineName + '_LZIkSpline' + '_surplus' + str(num_pma)+'_md'  
                    mc.createNode( 'multiplyDivide' , n = surplus_md )
                    mc.setAttr(str(surplus_md)+'.operation',2)
                    mc.connectAttr( str(plusMinusAverage)+'.output1D' , str(surplus_md)+'.input1X' )
                    mc.connectAttr( str(global_md) +'.outputX' , str(surplus_md)+'.input2X' )
                    
                    surplus_stretch_md = splineName + '_LZIkSpline' + '_surplusStretch' + str(num_pma)+'_md'  
                    mc.createNode( 'multiplyDivide' , n = surplus_stretch_md )
                    mc.setAttr(str(surplus_stretch_md)+'.operation',1)
                    
                    mc.connectAttr( str(surplus_md) +'.outputX' , str(surplus_stretch_md)+'.input1X' )
                    mc.connectAttr( str(allgrp)+'.stretch' , str(surplus_stretch_md)+'.input2X' )
                    
                    combinePMA = splineName + '_LZIkSpline' + '_combine' + str(num_pma)+'_pma'  
                    mc.createNode( 'plusMinusAverage' , n = combinePMA )
                    mc.connectAttr( str(surplus_stretch_md) +'.outputX' , str(combinePMA)+'.input1D[0]' )
                    mc.setAttr(  str(combinePMA)+'.input1D[1]' ,1 )
                    
                    combine_md_all.append(combinePMA)
                
                    mc.connectAttr( str(combinePMA)+'.output1D',str( seketon_all[num_pma] )+'.scaleX')
                    #----------------------------------------------------------------------------------
                    LateralScale_md = splineName + '_LZIkSpline' + '_lateral' + str(num_pma)+'_md'   
                    mc.createNode( 'multiplyDivide' , n = LateralScale_md )
                    
                    mc.connectAttr( str( seketon_all[num_pma] )+'.scaleX' ,  str(LateralScale_md)+'.input2Y' )
                    mc.connectAttr( str( seketon_all[num_pma] )+'.scaleX' ,  str(LateralScale_md)+'.input2Z' )
                    
                    mc.setAttr( str(LateralScale_md)+'.input1Y' , 1)
                    mc.setAttr( str(LateralScale_md)+'.input1Z' , 1)
                    mc.setAttr( str(LateralScale_md)+'.operation' , 2)
                    
                    mc.connectAttr( str(LateralScale_md)+'.outputY' , str( seketon_all[num_pma] )+'.scaleY' )
                    mc.connectAttr( str(LateralScale_md)+'.outputZ' , str( seketon_all[num_pma] )+'.scaleZ' )
                    
                    
                
                
                
                return allgrp
            else:
                pass            
    '''                
    def createSeketon(self,seketonInfo = [ ('' , (0,0,0) )] ,selectCancle = 1 ):
        for si in seketonInfo:
            if selectCancle==1:
                mc.select(cl=True)
            mc.joint( n = si[0] , p = si[1])
            parentObj = mc.listRelatives( si[0] ,p=True )
            print parentObj
            if parentObj == None:

                pass
                
            else:
                if mc.nodeType(parentObj[0]) == 'joint':

                    mc.joint( parentObj[0] ,e=True,zso=True,oj= 'xyz' , sao = 'yup' )
                else:
    '''
    def createSeketon(self,seketonInfo = [ ('' , (0,0,0) )] ,selectCancle = 1 ):
        
        
        if len(seketonInfo)==1:
            for si in seketonInfo:
                if selectCancle==1:
                    mc.select(cl=True)
                mc.joint( n = si[0] , p = si[1])
                parentObj = mc.listRelatives( si[0] ,p=True )
                #print parentObj
                if parentObj == None:
            
                    pass
                    
                else:
                    if mc.nodeType(parentObj[0]) == 'joint':
            
                        mc.joint( parentObj[0] ,e=True,zso=True,oj= 'xyz' , sao = 'yup' )
                    else:   
                        pass
        else:
            
            mc.select(cl=True)
            for si in seketonInfo:
   
                mc.joint( n = si[0] , p = si[1])
                parentObj = mc.listRelatives( si[0] ,p=True )
                #print parentObj
                if parentObj == None:
            
                    pass
                    
                else:
                    if mc.nodeType(parentObj[0]) == 'joint':
            
                        mc.joint( parentObj[0] ,e=True,zso=True,oj= 'xyz' , sao = 'yup' )
                    else:   
                        pass
            

